from PIL import Image, ImageDraw, ImageFont
import random, re, base64, io


# 生成随机验证码
def get_code(length):
    code = ''
    characters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'
    for _ in range(length):
        code += random.choice(characters)
    return code


# 生成验证码图片
def create_captcha(code: str, width, height, font_size, noise_points=100):
    image = Image.new('RGB', (width, height), (255, 255, 255))
    draw = ImageDraw.Draw(image)

    font = ImageFont.truetype('src/assets/font.ttf', font_size)  # 设置字体，需要提前安装字体到系统或指定字体文件路径
    x = (width - font_size * len(code)) / 2
    y = height / 2 - font_size
    i = 0
    for c in re.split(r"", code):
        if c != "":
            draw.text((x + font_size * i, y), text=c, font=font,
                      fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
            i = i + 1

    # 绘制噪点
    for _ in range(noise_points):
        x = random.randint(0, width - 1)
        y = random.randint(0, height - 1)
        draw.point((x, y), fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    # 绘制验证码文本
    # 保存图片
    # 创建一个BytesIO对象来暂存图片数据
    img_byte_arr = io.BytesIO()
    # 将图片保存到BytesIO对象中
    image.save(img_byte_arr, format='PNG')
    # 获取BytesIO对象中的图片数据并进行base64编码
    encoded_string = base64.b64encode(img_byte_arr.getvalue())
    # 将bytes类型转换成utf-8编码的字符串并返回
    return encoded_string.decode('utf-8')
